from typing import List, Tuple, Dict, Callable, Set
import multiprocessing as mp
from heapq import heappush, heappop
from itertools import combinations
from copy import deepcopy
import numpy as np
from queue import Queue

# ST Astar算法的规划部分
from ..low_level.AStar.planner import Planner as STPlanner

from .constraint_tree import CTNode
from .constraints import Constraints
from .agent import Agent

class Planner:

    def __init__(self, grid_size: int,
                       robot_radius: int,
                       static_obstacles: List[Tuple[int, int]]):

        self.robot_radius = robot_radius
        self.st_planner = STPlanner(grid_size, robot_radius, static_obstacles) 

        self.edge_collision_times = 0

    def plan(self, starts: List[Tuple[int, int]],
                   goals: List[Tuple[int, int]],
                   max_iter:int = 200,
                   low_level_max_iter:int = 100,
                   max_process:int = 10,
                   debug:bool = False) -> np.ndarray:

        self.debug = debug
        self.low_level_max_iter = low_level_max_iter



        # Do goal assignment
        # 分派终点，起点。
        self.agents = []
        for i in range(len(starts)):
            self.agents.append(Agent(starts[i], goals[i]))
        # 约束类
        self.constraints = Constraints()

        # Compute path for each agent using low level planner
        solution = dict((agent, self.calculate_path(agent, self.constraints, None)) for agent in self.agents)

        # 添加 Batch-path ,还需要完善，现在只是最基础的运算推出群路径，
        # 而不是真实的通过A*携带不同的启发算法算出来的
        batch_path = {}
        for path in solution.values():
            b_path = self.batch_path(path, 80)
            batch_path.update(b_path)


        open = []
        # 所有path都不为零，即每个车均有解
        if all(len(path) != 0 for path in solution.values()):
            # Make root node
            node = CTNode(self.constraints, solution, batch_path)
            # Min heap for quick extraction
            open.append(node)

        # manager = mp.Manager()
        iter_ = 0
        while open and iter_ < max_iter:
            iter_ += 1

            # results = manager.list([])
            # processes = []
            # # Default to 10 processes maximum
            # for _ in range(max_process if len(open) > max_process else len(open)):
            #     p = mp.Process(target=self.search_node, args=[heappop(open), results])
            #     processes.append(p)
            #     p.start()
            # for p in processes:
            #     p.join()
            # for result in results:
            #     if len(result) == 1:
            #         if debug:
            #             print('CBS_MAPF: Paths found after about {0} iterations'.format(4 * iter_))
            #         return result[0]
            #     if result[0]:
            #         heappush(open, result[0])
            #     if result[1]:
            #         heappush(open, result[1])

            results = []
            self.search_node(heappop(open), results, batch_path)
            for result in results:
                if len(result) == 1:
                    if debug: 
                        print('CBS_MAPF: Paths found after about {0} iterations'.format(4 * iter_))
                    return result[0]
                if result[0]:
                    heappush(open, result[0])
                if result[1]:
                    heappush(open, result[1])

        if debug:
            print('CBS-MAPF: Open set is empty, no paths found.')
        return np.array([])

    '''
    Abstracted away the cbs search for multiprocessing.
    The parameters open and results MUST BE of type ListProxy to ensure synchronization.
    '''
    def search_node(self, best: CTNode, results, batch_path):
        agent_i, agent_j, time_of_conflict = self.validate_paths(self.agents, best)
        
        # If there is not conflict, validate_paths returns (None, None, -1)
        if agent_i is None:
            results.append((self.reformat(self.agents, best.solution),))
            return
        # Calculate new constraints
        agent_i_constraint = self.calculate_constraints(best, agent_i, agent_j, time_of_conflict)
        agent_j_constraint = self.calculate_constraints(best, agent_j, agent_i, time_of_conflict)

        # Calculate new paths
        agent_i_path = self.calculate_path(agent_i,
                                           agent_i_constraint,
                                           self.calculate_goal_times(best, agent_i, self.agents))
        agent_j_path = self.calculate_path(agent_j,
                                           agent_j_constraint,
                                           self.calculate_goal_times(best, agent_j, self.agents))

        # Replace old paths with new ones in solution
        solution_i = best.solution
        solution_j = deepcopy(best.solution)
        solution_i[agent_i] = agent_i_path
        solution_j[agent_j] = agent_j_path

        node_i = None
        if all(len(path) != 0 for path in solution_i.values()):
            node_i = CTNode(agent_i_constraint, solution_i, batch_path)

        node_j = None
        if all(len(path) != 0 for path in solution_j.values()):
            node_j = CTNode(agent_j_constraint, solution_j, batch_path)

        results.append((node_i, node_j))


    '''
    Pair of agent, point of conflict
    '''
    def validate_paths(self, agents, node: CTNode):
        # Check collision pair-wise
        for agent_i, agent_j in combinations(agents, 2):
            time_of_conflict = self.vertex_collision(node.solution, agent_i, agent_j)
            # time_of_conflict=-1 if there is not conflict
            if time_of_conflict == -1:
                time_of_conflict = self.edge_collision(node.solution, agent_i, agent_j)
                if time_of_conflict == -1:
                    continue
            return agent_i, agent_j, time_of_conflict
        return None, None, -1


    # def identify_conflict(self, node, agents):
        # for agent_i, agent_j in combinations(agents, 2):
                    
    
    def new_vertex_collision(self, solution, agent_i, agent_j, batch_path) -> int:
        conflict = 0
        con_list = []
        idx_list = []
        for idx, (point_i, point_j) in enumerate(zip(solution[agent_i], solution[agent_j])):
            if self.dist(point_i, point_j) < 2*self.robot_radius:
                conflict = 1
                con_list.append(conflict)
                idx_list.append(idx)
                break
        if conflict == 0:
            con_list.append(conflict)
            return -1
        else:
            con_list_1, idx_list_1 = self.oneCompareMulti(solution, agent_i, agent_j, batch_path)
            con_list_itoj = con_list + con_list_1

            con_list_2, idx_list_2 = self.oneCompareMulti(solution, agent_j, agent_i, batch_path)
            con_list_jtoi = con_list + con_list_2
        

        return conflict

    def oneCompareMulti(self, solution, agent_i, agent_j, batch_path):
        conflict = 0
        con_list = []
        idx_list = []
        multi_path = batch_path[tuple(agent_j.start)]
        for path in multi_path:
            for idx, (point_i, point_j) in enumerate(zip(solution[agent_i], path)):
                if self.dist(point_i, point_j) < 2*self.robot_radius:
                    conflict = 1
                    con_list.append(conflict)
                    idx_list.append(idx)
                    break
            if con_list == 0:
                con_list.append(conflict)
                idx_list.append(0)

        return con_list, idx_list

        

    def vertex_collision(self, solution: Dict[Agent, np.ndarray], agent_i: Agent, agent_j: Agent) -> int:
        for idx, (point_i, point_j) in enumerate(zip(solution[agent_i], solution[agent_j])):
            if self.dist(point_i, point_j) > 2*self.robot_radius:
                continue
            tangle_collision = self.tangle_collision(solution, agent_i, agent_j)
            if tangle_collision:
                return 'tangle_collision'
            return idx
        return -1

    # Check edge_collision
    def edge_collision(self, solution, agent_i, agent_j):

        frequency = 0

        queue = []
        queue.append((solution[agent_i][0], solution[agent_j][0]))
        queue.append((solution[agent_i][1], solution[agent_j][1]))

        idx_list = []
        points_list = []
        for idx, (point_i, point_j) in enumerate(zip(solution[agent_i], solution[agent_j])):
            idx_list.append(idx)
            points_list.append((point_i, point_j))

        # idx_list.insert(0,0)
        # idx_list.insert(0,0)

        points_list.pop(0)
        points_list.pop(0)
        points_list.append((solution[agent_i][-1], solution[agent_j][-1]))

        for idx, (point_i, point_j) in zip(idx_list, points_list):

            if len(queue) == 2:

                first = queue[0]
                follow = queue[1]

                if not ((first[1] - follow[0]).any() or (first[0] - follow[1]).any()):
                    frequency += 1
                    return idx-1, idx, idx+1

                # if not ((first[1] - last[0]).any() and (first[0] == last[1]).any()):
                #     return idx-3, idx-2, idx-1

                queue.pop(0)


            queue.append((point_i, point_j))
        return -1
            
    def tangle_collision(self, solution, agent_i, agent_j):
        # 定义两车起点终点x,y坐标
        S_i = solution[agent_i][0]
        S_j = solution[agent_j][0]
        S_i_x = S_i[0]
        S_i_y = S_i[1]
        S_j_x = S_j[0]
        S_j_y = S_j[1]

        G_i = solution[agent_i][-1]
        G_j = solution[agent_j][-1]
        G_i_x = G_i[0]
        G_i_y = G_i[1]
        G_j_x = G_j[0]
        G_j_y = G_j[1]
        # 确保两车当前路径是满足曼哈顿距离的
        i_manhattan =(abs(S_i_x-G_i_x)/80+abs(S_i_y-G_i_y)/80+1==len(solution[agent_i])) 
        j_manhattan =(abs(S_j_x-G_j_x)/80+abs(S_j_y-G_j_y)/80+1==len(solution[agent_j]))
        # 确保两车起点到终点，在x与y方向上共向
        self_x_tangle = ((S_i_x-G_i_x)*(S_j_x-G_j_x) >= 0)
        self_y_tangle = ((S_i_y-G_i_y)*(S_j_y-G_j_y) >= 0)
        # 确保两车起点与终点的可行路径是纠缠的
        x_tangle = ((S_i_x-S_j_x)*(G_i_x-G_j_x) <= 0)
        y_tangle = ((S_i_y-S_j_y)*(G_i_y-G_j_y) <= 0)

        tangle_collision = all([i_manhattan, j_manhattan, self_x_tangle, self_y_tangle, x_tangle, y_tangle])
        return tangle_collision

    @staticmethod
    def dist(point1: np.ndarray, point2: np.ndarray) -> int:
        return int(np.linalg.norm(point1-point2, 2))  # L2 norm

    def calculate_constraints(self, node: CTNode,
                                    constrained_agent: Agent,
                                    unchanged_agent: Agent,
                                    time_of_conflict) -> Constraints:


        contrained_path = node.solution[constrained_agent]
        unchanged_path = node.solution[unchanged_agent]
        # 一般约束
        if type(time_of_conflict) == int:
            pivot = unchanged_path[time_of_conflict]
            conflict_end_time = time_of_conflict
            try:
                while self.dist(contrained_path[conflict_end_time], pivot) < 2*self.robot_radius:
                    conflict_end_time += 1
            except IndexError:
                pass
            return node.constraints.fork(constrained_agent, tuple(pivot.tolist()), time_of_conflict, conflict_end_time)
        # 边碰撞约束
        if len(time_of_conflict) == 3:
            pivot =  unchanged_path[time_of_conflict[1]]

            conflict_end_time = time_of_conflict[2]
            # 检测到边碰撞的次数
            # self.edge_collision_times += 1
            
            # if self.edge_collision_times == 1:
            #     conflict_end_time = time_of_conflict[1]
            # else:
            #     conflict_end_time = time_of_conflict[2]

            return node.constraints.fork_edge(constrained_agent, 
                                            pivot,
                                            time_of_conflict[1], conflict_end_time)
        # 纠缠碰撞约束
        if type(time_of_conflict) != int:
            pivot = contrained_path[0]
            conflict_end_time = 2

            return node.constraints.fork_tangle(constrained_agent, 
                                            pivot,
                                            1, conflict_end_time)            

    def calculate_goal_times(self, node: CTNode, agent: Agent, agents: List[Agent]):
        solution = node.solution
        goal_times = dict()
        for other_agent in agents:
            if other_agent == agent:
                continue
            time = len(solution[other_agent]) - 1
            goal_times.setdefault(time, set()).add(tuple(solution[other_agent][time]))
        return goal_times

    '''
    Calculate the paths for all agents with space-time constraints
    '''
    def calculate_path(self, agent: Agent, 
                       constraints: Constraints, 
                       goal_times: Dict[int, Set[Tuple[int, int]]]) -> np.ndarray:
        return self.st_planner.plan(agent.start, 
                                    agent.goal, 
                                    constraints.setdefault(agent, dict()), 
                                    semi_dynamic_obstacles=goal_times,
                                    max_iter=self.low_level_max_iter, 
                                    debug=self.debug)

    '''
    Reformat the solution to a numpy array
    '''
    @staticmethod
    def reformat(agents: List[Agent], solution: Dict[Agent, np.ndarray]):
        
        solution = Planner.pad(solution)
        reformatted_solution = []
        for agent in agents:
            reformatted_solution.append(solution[agent])
        return np.array(reformatted_solution)

    '''
    Pad paths to equal length, inefficient but well..
    '''
    @staticmethod
    def pad(solution: Dict[Agent, np.ndarray]):
        max_ = max(len(path) for path in solution.values())
        for agent, path in solution.items():
            if len(path) == max_:
                continue
            padded = np.concatenate([path, np.array(list([path[-1]])*(max_-len(path)))])
            solution[agent] = padded
        return solution

    def batch_path(self, path, grid_size):
        
        start = path[0]
        goal = path[-1]

        start_x = path[0][0]
        start_y = path[0][1]
        goal_x = path[-1][0]
        goal_y = path[-1][1]

        # min_x = min(start_x, goal_x)
        # min_y = min(start_y, goal_y)
        # max_x = max(start_x, goal_x)
        # max_y = max(start_y, goal_y)

        if start_x <= goal_x:
            x_value = self.creat_list(start_x, goal_x, grid_size)
        else:
            x_value = self.creat_list(goal_x, start_x, grid_size)
            x_value = sorted(x_value, reverse=True)

        if start_y <= goal_y:
            y_value = self.creat_list(start_y, goal_y, grid_size)
        else:
            y_value = self.creat_list(goal_y, start_y, grid_size)
            y_value = sorted(y_value, reverse=True)
        
        # 先在X方向运动到底的路径
        x_path = []
        for x in x_value:
            x_path.append((x,start_y))
        
        x_path.pop()
        for y in y_value:
            x_path.append((goal_x,y))
        
        x_path = np.array(x_path)

        y_path = []
        for y in y_value:
            y_path.append((start_x,y))

        y_path.pop()
        for x in x_value:
            y_path.append((x,goal_y))

        y_path = np.array(y_path)

        batch_path = np.array([x_path,y_path])

        tup_start = tuple(start)

        path_dict = dict([(tup_start, batch_path)]) 

        return path_dict

    def creat_list(self, r1, r2, grid_size):
        return [item for item in range(r1, r2 + 1, grid_size)]